#include "stdafx.h"
#include "Renderer.h"
#include "RenderBuffer.h"
#include "SobolSequenceGenerator.h"
#include <intrin.h>
#include <map>

namespace RayTracePro
{
	using namespace Concurrency;
	struct Point
	{
		int x, y;
		bool operator < (const Point & p) const
		{
			if (x < p.x)
				return true;
			else if (x == p.x)
				return y < p.y;
			else
				return false;
		}
		bool operator == (const Point & p) const
		{
			return p.x == x && p.y == y;
		}
	};

	class Hilbert
	{
		GxWinSystem::Threading::gxLock lock;
		// code from wikipedia
		//convert d to (x,y)
		static void d2xy(int n, int d, int *x, int *y)
		{
			int rx, ry, s, t=d;
			*x = *y = 0;
			for (s=1; s<n; s*=2)
			{
				rx = 1 & (t/2);
				ry = 1 & (t ^ rx);
				rot(s, x, y, rx, ry);
				*x += s * rx;
				*y += s * ry;
				t /= 4;
			}
		}
 
		//rotate/flip a quadrant appropriately
		static void rot(int n, int *x, int *y, int rx, int ry)
		{
			int t;
			if (ry == 0)
			{
				if (rx == 1)
				{
					*x = n-1 - *x;
					*y = n-1 - *y;
				}
				t  = *x;
				*x = *y;
				*y = t;
			}
		}
	private:
		std::map<Point, RefPtr<Array<Point, true>>> blocks;
	public:
		// size must be power of 2
		void FillSquare(int offsetX, int offsetY, int size, Array<Point, true> & curve)
		{
			Point p;
			p.x = size; p.y = size;
			auto block = blocks.find(p);
			if (block == blocks.end())
			{
				RefPtr<Array<Point, true>> blk = new Array<Point, true>();
				blk->SetSize(size * size);
				for (int i = 0; i<size*size; i++)
				{
					Point p;
					int ix, iy;
					d2xy(size, i, &ix, &iy);
					p.x = ix;
					p.y = iy;
					blk->GetItem(i) = p;
					p.x += offsetX;
					p.y += offsetY;
					curve.Add(p);
				}
				blocks[p] = blk;
			}
			else
			{
				auto blk = block->second;
				for (int i = 0; i<blk->Count(); i++)
				{
					Point p = blk->GetItem(i);
					p.x += offsetX;
					p.y += offsetY;
					curve.Add(p);
				}
			}
		}

		void FillBlock(int offsetX, int offsetY, int width, int height, Array<Point, true> & curve)
		{
			if (width > 4 && height > 4)
			{
				int size = min(width, height);
				int n = 1;
				while ((1<<n) < size)
					n++;
				n--;
				//if (n>5) n = 5;
				int blockSize = 1<<n;
				FillSquare(offsetX, offsetY, blockSize, curve);
				if (width > height)
				{
					FillBlock(offsetX + blockSize, offsetY, width-blockSize, height, curve);
					FillBlock(offsetX, offsetY + blockSize, blockSize, height-blockSize, curve);
				}
				else
				{
					FillBlock(offsetX, offsetY + blockSize, width, height-blockSize, curve);
					FillBlock(offsetX + blockSize, offsetY, width-blockSize, blockSize, curve);
					
				}
			}
			else
			{
				for (int i = 0; i<width; i++)
					for (int j = 0; j<height; j++)
					{
						Point p;
						p.x = i + offsetX;
						p.y = j + offsetY;
						curve.Add(p);
					}
			}
		}

		Array<Point, true> * FillBlock(int width, int height)
		{
			Point p;
			p.x = width;
			p.y = height;
			//lock.Lock();
			auto block = blocks.find(p);
			//lock.Unlock();
			if (block == blocks.end() || block->first.x != width || block->first.y != height)
			{
				RefPtr<Array<Point, true>> cache = new Array<Point, true>();
				cache->SetCapacity(width * height);
				FillBlock(0,0, width, height, *cache);
				//lock.Lock();
				blocks[p] = cache;
				//lock.Unlock();
				return cache.operator->();
			}
			else
			{
				return block->second.operator->();
			}
		}
	};

	ShadingBuffer::ShadingBuffer(RenderContext * renderContext, int allocId, DeviceList<float> * buffer)
	{
		context = renderContext;
		allocationId = allocId;
		buffer = buffer;
	}
	void ShadingBuffer::Free()
	{
		if (context)
			context->FreeShadingBuffer(allocationId);
	}

	Renderer::Renderer(RayTraceSystem * system)
	{
		completedTasks = 0;
		hilbert = new Hilbert();
		this->system = system;
		this->kernelComputeCameraRays = GpuDevice::CreateKernel("GenerateCameraRays");
		this->kernelDebugShade = GpuDevice::CreateKernel("DefaultShade");
	}
	Renderer::~Renderer()
	{
		delete hilbert;
	}
	void Renderer::SetScene(Scene * scene)
	{
		this->scene = scene;
	}
	void Renderer::SetScreenBuffer(RenderBuffer * buf)
	{
		buffer = buf;
		const int MaxRaysPerBlock = 512*512*2;
		int blockSize = Align((int)(sqrt(MaxRaysPerBlock/(float)system->AntiAliasSamples)), 8);
		int bWidth = (buf->Width + blockSize - 1)/blockSize;
		int bHeight = (buf->Height + blockSize - 1)/blockSize;
		// average blocks:
		if (bWidth > 1 || bHeight > 1)
		{
			blockSize = min(buf->Width/bWidth, buf->Height/bHeight);
			bWidth = (buf->Width + blockSize - 1)/blockSize;
			bHeight = (buf->Height + blockSize - 1)/blockSize;
		}
		tasks.SetCapacity(bWidth * bHeight);
		tasks.Clear();
		for (int i = 0; i<bWidth; i++)
			for (int j = 0; j<bHeight; j++)
			{
				RenderTask t;
				t.X = i * blockSize;
				t.Y = j * blockSize;
				t.Width = min(buf->Width - t.X, blockSize);
				t.Height = min(buf->Height - t.Y, blockSize);
				hilbert->FillBlock(t.Width, t.Height);
				tasks.Add(t);
			}
		completedTasks = tasks.Count();
	}

	void Renderer::DefaultShade(bool GPU, DeviceList<float> & colors, DeviceList<Ray> & rays, DeviceList<RayDifferential> & rayDiffs, DeviceList<int> & rayIds, int start, int end, RenderContext & context)
	{
		colors.SetSize(rays.Count()*4);
		memset(colors.Buffer(), 0, sizeof(Vec4)*rays.Count());
	}

	void Renderer::DebugTestShade(bool gpu, DeviceList<Vec4> & colors, DeviceList<DifferentialGeometry> & intersections)
	{
		if (gpu)
		{
			GpuDevice::InvokeKernel(kernelDebugShade, Align(colors.Count(), 256), 256)
				(colors.Count(), intersections.DeviceBuffer(false), colors.DeviceBuffer(false));
			colors.Synchronize();
		}
		else
		{
			Vec3 dir(-1.0f, -0.7f, -1.0f);
			Vec3::Normalize(dir, dir);
			parallel_for (0, intersections.Count(), [&](const int i)
			{
				Vec4::Scale(colors[i], Vec4(1.0f, 1.0f, 1.0f, 1.0f), max(0.0f, Vec3::Dot(intersections[i].Normal, dir)));
				colors[i].w = 1.0f;
			});
		}
		
	}
	
	void Renderer::CpuRenderThread()
	{
		system->InitRenderThread();
		//try
		{
			system->StartBenchmark("CPU Render");
			while (true)
			{
				long taskId = _InterlockedIncrement(&taskPtr) - 1;
				if (taskId < tasks.Count())
				{
					RenderTask t = tasks[taskId];
					RenderBlock(false, t.X, t.Y, t.Width, t.Height, cpuContext);
					_InterlockedIncrement(&completedTasks);
					taskCompletedByCPU++;
					OnBlockFinished(t);
				}
				else
					break;
			}
			system->StopBenchmark();
		}
	/*	catch(...)
		{
			printf("Unknown exception in render thread.\n");
		}*/
		system->FinalizeRenderThread();
	}

	void Renderer::GpuRenderThread()
	{
		system->InitRenderThread();
		try
		{
			system->StartBenchmark("GPU Render");
			while (true)
			{
				long taskId = _InterlockedIncrement(&taskPtr) - 1;
				if (taskId < tasks.Count())
				{
					RenderTask t = tasks[taskId];
					RenderBlock(true, t.X, t.Y, t.Width, t.Height, gpuContext);
					_InterlockedIncrement(&completedTasks);
					OnBlockFinished(t);
				}
				else
					break;
			}
			system->StopBenchmark();
		}
		catch(cl::Error err)
		{
			printf("OpenCL error %d in %s\n", err.err(), err.what());
		}
		catch(...)
		{
			printf("Unknown exception in render thread.\n");
		}
		system->FinalizeRenderThread();
	}

	void Renderer::Cancel()
	{
		cancelled = true;
		while (completedTasks < tasks.Count())
			Sleep(10);
		RenderFinished.Invoke(false, RenderStatistics());
	}

	bool Renderer::RenderFrameAsync(RenderDevice device)
	{
		if (completedTasks < tasks.Count())
			return false;
		QueryPerformanceCounter(&startTime);
		taskPtr = 0;
		completedTasks = 0;
		taskCompletedByCPU = 0;
		cancelled = false;
		if (device == CPU || device == Both)
		{
			QueryPerformanceCounter(&startTime);
			GxWinSystem::Threading::gxThreadProc *cpuThreadProc = new GxWinSystem::Threading::gxThreadProc();
			cpuThreadProc->Bind(this, &Renderer::CpuRenderThread);
			cpuThread = new GxWinSystem::Threading::gxThread(cpuThreadProc);
			cpuThread->Start();
		}
		if (device == GPU || device == Both)
		{
			GxWinSystem::Threading::gxThreadProc *gpuThreadProc = new GxWinSystem::Threading::gxThreadProc();
			gpuThreadProc->Bind(this, &Renderer::GpuRenderThread);
			gpuThread = new GxWinSystem::Threading::gxThread(gpuThreadProc);
			gpuThread->Start();
		}
		return true;
	}

	RenderContext::RenderContext()
	{
		RandX.SetSize(1024);
		RandY.SetSize(1024);
		SobolGenerator randGen(1334, 5982, 2045);
		for (int i = 0; i<1024; i++)
		{
			RandX[i] = (float)randGen.NextX();
			RandY[i] = (float)randGen.NextY();
		}
		for (int i = 0; i<MaxShadingBuffers; i++)
		{
			shadingBufferAllocationState[i] = false;
		}
	}

	void Renderer::RenderBlock(bool GPU, int x0, int y0, int width, int height, RenderContext & context)
	{
		if (GPU)
			system->StartBenchmark("RenderBlock(GPU)");
		else
			system->StartBenchmark("RenderBlock");
		system->StartBenchmark("GetRenderSequence");
		Array<Point, true> * sequence = hilbert->FillBlock(width, height);
		DeviceList<Vec2> &offsetSequence = context.PixelPositions;
		offsetSequence.SetSize(width * height * system->AntiAliasSamples);
		int aaSamples = system->AntiAliasSamples;
		parallel_for (0, sequence->Count(), [&](int i)
		{
			if (aaSamples > 1)
			{
				for (int k = 0; k < aaSamples; k++)
				{
					offsetSequence[i * aaSamples + k].x = (float)(sequence->GetItem(i).x + x0 + context.RandX[(k+i*3)&1023] - 0.5);
					offsetSequence[i * aaSamples + k].y = (float)(sequence->GetItem(i).y + y0 + context.RandY[(k+i*3)&1023] - 0.5);
				}
			}
			else
			{
				offsetSequence[i].x = (float)(sequence->GetItem(i).x + x0);
				offsetSequence[i].y = (float)(sequence->GetItem(i).y + y0);
			}
		});
		system->StopBenchmark();
		system->StartBenchmark("ComputeCameraRays");
		DeviceList<Ray> & rays = context.Rays;
		DeviceList<RayDifferential> & rayDiffs = context.RayDiffs;
		DeviceList<DifferentialGeometry> & intersections = context.Intersections;
		int rayCount = width * height * system->AntiAliasSamples;
		intersections.SetSize(rayCount);
		context.IntermediateIntersections.SetSize(rayCount);
		rays.SetSize(rayCount);
		rayDiffs.SetSize(rayCount);
		int alignedSize = (int)Ceil(rayCount/256.0f)*256;
		if (GPU)
		{
			try
			{
				GpuDevice::InvokeKernel(kernelComputeCameraRays, alignedSize, 256)
					(rayCount, Camera.GetRight().x, Camera.GetUp().y, Camera.GetCenterX(), Camera.GetCenterY(),
					Camera.GetScreenZ(), Camera.CameraTransform, offsetSequence.DeviceBuffer(), 
					rays.DeviceBuffer(false), rayDiffs.DeviceBuffer(false));//.Wait();
			}
			catch(cl::Error err)
			{
				printf("OpenCL kernel launch error %d\n", err.err());
			}	
		}
		else
		{
			parallel_for (0, rays.Count(), [&](int id)
			{	
				Ray ray;
				RayDifferential rayDiff;
				auto p = offsetSequence[id];
				rays[id].RayId = id;
				CameraProjection::GenerateRay(Camera, rays[id], rayDiffs[id], p.x, p.y);
			});
		}
		system->StopBenchmark();
		system->StartBenchmark("TraceRays");
		Vec3 camPos(Camera.CameraTransform.values[12],Camera.CameraTransform.values[13],Camera.CameraTransform.values[14]);
		if (GPU)
		{
			scene->TraceRay_BVH_GPU(intersections, rays, rayDiffs, context.TraverseContinuation, context.HeadFlags, context.ScanHeadFlags, context.Groups, context.skeleton, context.IntermediateIntersections);
		}
		else
			scene->TraceRay_BVH(intersections, rays, rayDiffs, context.TraverseContinuation, context.HeadFlags, context.ScanHeadFlags, context.Groups, context.skeleton);
		system->StopBenchmark();
		system->StartBenchmark("Shading");
		DeviceList<Vec4> & colors = context.Colors;
		colors.SetSize(intersections.Count());
		DebugTestShade(GPU, colors, intersections);
		system->StopBenchmark();
		system->StartBenchmark("FillFrameBuffer");
		float avgWeight = 1.0f/system->AntiAliasSamples;
		parallel_for (0, width* height, [&](int i)
		{
			Point & p = sequence->GetItem(i);
			Vec4 color; color.SetZero();
			for (int j = 0; j<aaSamples; j++)
				color += colors[i*aaSamples + j];
			color *= avgWeight;
			buffer->SetPixel(p.x + x0, p.y + y0, color);
		});
		system->StopBenchmark();
		system->StopBenchmark();
	}

	void Renderer::Shade(bool GPU,
						 DeviceList<float> & colors,
						 DeviceList<DifferentialGeometry> & intersections, 
						 DeviceList<Ray> & rays, 
						 DeviceList<RayDifferential> & rayDiffs, 
						 RenderContext & context)
	{
		if (intersections.Count() == 0)
			return;
		colors.SetSize(rays.Count()*4);
		if (GPU)
			throw "GPU Shading Not Implemented";
		// Sort intersections by MaterialId
		DeviceList<int> & headFlags = context.HeadFlags;
		DeviceList<int> & scanHeadFlags = context.ScanHeadFlags;
		DeviceList<int> skeleton;// = context.skeleton;
		DeviceList<SortGroup> & groups = context.Groups;
		headFlags.SetSize(intersections.Count());
		scanHeadFlags.SetSize(intersections.Count());
		skeleton.SetSize(intersections.Count());
		headFlags[0] = 1;
		parallel_for(1, intersections.Count(), [&](int i)
		{
			int head = (intersections[i-1].MaterialId != intersections[i].MaterialId);
			headFlags[i] = head;
			if (head)
				scanHeadFlags[i] = i;
			else
			{
				scanHeadFlags[i] = 1;
			}
		});
		int groupCount = ExclusiveScan(scanHeadFlags);
		groups.SetSize(groupCount);
		parallel_for(0, intersections.Count(), [&](int i)
		{
			if (headFlags[i] == 1)
			{
				groups[scanHeadFlags[i]].Base = i;
				groups[scanHeadFlags[i]].GroupId = intersections[i].MaterialId;
			}
		});
		parallel_for(0, groups.Count(), [&](int i)
		{
			if (i<groups.Count()-1)
				groups[i].Count = groups[i+1].Base-groups[i].Base;
			else
			{
				groups[i].Count = intersections.Count()-groups[i].Base;
			}
		});
		parallel_radixsort(groups.Buffer(), groups.Buffer()+groups.Count(), [](const SortGroup &g){return g.GroupId;});
		// build skeleton
		int ptr = 0;
		for (int i = 0; i<groups.Count(); i++)
		{
			int base = groups[i].Base;
			for (int j = 0; j<groups[i].Count; j++)
			{
				skeleton[ptr++] = base + j;
			}
		}
		int gPtr = 0;
		int start = 0;
		// foreach group
		while (gPtr < groups.Count())
		{
			int end = start + groups[gPtr].Count;
			while (gPtr + 1 < groups.Count() && groups[gPtr+1].GroupId == groups[gPtr].GroupId)
			{
				gPtr++;
				end += groups[gPtr].Count;
			}
			// Shade skeleton[start, end]
			if (groups[gPtr].GroupId != -1)
			{
				int materialId = groups[gPtr].GroupId;
				//scene->Materials[materialId]->Shade(colors, intersections, rays, rayDiffs, skeleton, start, end, context);
			}
			else
			{
				// no intersections, shade environment
				DefaultShade(GPU, colors, rays, rayDiffs, skeleton, start, end, context);
			}
			start = end;
			gPtr++;
		}
	}
}